package com.alison.datastream.chapter4_join;

import org.apache.flink.api.common.functions.CoGroupFunction;
import org.apache.flink.api.java.tuple.Tuple2;
import org.apache.flink.streaming.api.TimeCharacteristic;
import org.apache.flink.streaming.api.datastream.DataStream;
import org.apache.flink.streaming.api.environment.StreamExecutionEnvironment;
import org.apache.flink.streaming.api.functions.source.SourceFunction;
import org.apache.flink.streaming.api.windowing.assigners.TumblingProcessingTimeWindows;
import org.apache.flink.streaming.api.windowing.time.Time;
import org.apache.flink.util.Collector;

import java.util.concurrent.TimeUnit;

public class F3_CoGroupExample {

    // CoGroup算子：将两个数据流按照key进行group分组，并将数据流按key进行分区的处理，最终合成一个数据流（与join有区别，不管key有没有关联上,最终都会合并成一个数据流）
    public static void main(String[] args) throws Exception {

        StreamExecutionEnvironment env = StreamExecutionEnvironment.getExecutionEnvironment();
        // 使用EventTime时间语义
        env.setParallelism(1);
        // 如果配置了event_time，就配置watermark
        env.setStreamTimeCharacteristic(TimeCharacteristic.ProcessingTime);

        DataStream<Tuple2<String, Integer>> input1 = env.addSource(new MySource(1));
        DataStream<Tuple2<String, Integer>> input2 = env.addSource(new MySource(2));
        DataStream<String> coGroupResult = input1
                .coGroup(input2)
                .where(i1 -> i1.f0)// input1.key
                .equalTo(i2 -> i2.f0)// input2.key ,  where input1.key=input2.key
                .window(TumblingProcessingTimeWindows.of(Time.seconds(10)))
                .apply(new MyCoGroupFunction());

        coGroupResult.print();
        env.execute("window cogroup function");
    }

    /**
     * IN1: Tuple2<String, Integer>
     * IN2: Tuple2<String, Integer>
     * OUT: String
     */
    public static class MyCoGroupFunction implements CoGroupFunction<Tuple2<String, Integer>, Tuple2<String, Integer>, String> {
        @Override
        public void coGroup(Iterable<Tuple2<String, Integer>> input1, Iterable<Tuple2<String, Integer>> input2, Collector<String> out) {
            input1.forEach(element -> System.out.println("input1 :" + element.f1));
            input2.forEach(element -> System.out.println("input2 :" + element.f1));
        }
    }

    public static class MySource implements SourceFunction<Tuple2<String, Integer>> {
        private volatile boolean isRunning = true;

        int count = 1;

        public MySource(int count) {
            this.count = count;
        }

        @Override
        public void run(SourceContext ctx) throws Exception {

            while (isRunning) {
                Tuple2<String, Integer> data = Tuple2.of(String.valueOf(count), count);
                ctx.collect(data);
                count++;
                TimeUnit.SECONDS.sleep(2);
            }
        }

        @Override
        public void cancel() {
            isRunning = false;
        }
    }
}
